{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, '../../../Utilities/')\n",
    "import argparse\n",
    "import os\n",
    "import torch\n",
    "from collections import OrderedDict\n",
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.io\n",
    "from scipy.interpolate import griddata\n",
    "from plotting import newfig, savefig\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "import matplotlib.gridspec as gridspec\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.utils import save_image\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "from torchvision import datasets\n",
    "from torch.autograd import Variable\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "import seaborn as sns\n",
    "import pylab as py\n",
    "import time\n",
    "from pyDOE import lhs\n",
    "import warnings\n",
    "sys.path.insert(0, '../../../Scripts/')\n",
    "from models_pde import Generator, Discriminator, Q_Net\n",
    "from pig import *\n",
    "# from ../Scripts/helper import *\n",
    "\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "np.random.seed(1234)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CUDA support \n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device('cuda:0')\n",
    "else:\n",
    "    device = torch.device('cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyper-parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 30000\n",
    "lambda_phy = 1\n",
    "lambda_q = 0.5\n",
    "\n",
    "noise = 0.1\n",
    "\n",
    "#architecture for the models\n",
    "d_hid_dim = 50 \n",
    "d_num_layer = 2\n",
    "\n",
    "g_hid_dim = 50\n",
    "g_num_layer = 4\n",
    "\n",
    "q_hid_dim = 50\n",
    "q_num_layer = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_u = 100\n",
    "N_i = 50\n",
    "N_f = 10000\n",
    "data = scipy.io.loadmat('../../../datasets/burgers_shock.mat')\n",
    "\n",
    "t = data['t'].flatten()[:,None]\n",
    "x = data['x'].flatten()[:,None]\n",
    "Exact = np.real(data['usol']).T\n",
    "\n",
    "X, T = np.meshgrid(x,t)\n",
    "\n",
    "\n",
    "X_star = np.hstack((X.flatten()[:,None], T.flatten()[:,None]))\n",
    "u_star = Exact.flatten()[:,None] \n",
    "\n",
    "# Doman bounds\n",
    "lb = X_star.min(0)\n",
    "ub = X_star.max(0)\n",
    "\n",
    "# initial conditions t = 0\n",
    "xx1 = np.hstack((X[0:1,:].T, T[0:1,:].T))\n",
    "uu1 = Exact[0:1,:].T\n",
    "\n",
    "# boundary conditions x = lb\n",
    "xx2 = np.hstack((X[:,0:1], T[:,0:1]))\n",
    "uu2 = Exact[:,0:1]\n",
    "\n",
    "# boundary conditions, x = ub\n",
    "xx3 = np.hstack((X[:,-1:], T[:,-1:]))\n",
    "uu3 = Exact[:,-1:]\n",
    "\n",
    "X_u_train = np.vstack([xx2, xx3]) \n",
    "u_train = np.vstack([uu2, uu3])\n",
    "\n",
    "X_f_train = lb + (ub-lb)*lhs(2, N_f)\n",
    "X_f_train = np.vstack([X_f_train, X_u_train, xx1])\n",
    "\n",
    "# selecting N_u boundary points for training\n",
    "idx = np.random.choice(X_u_train.shape[0], N_u, replace=False)\n",
    "X_u_train = X_u_train[idx, :]\n",
    "u_train = u_train[idx,:]\n",
    "\n",
    "# selecting N_i initial points for training\n",
    "idx = np.random.choice(xx1.shape[0], N_i, replace=False)\n",
    "X_i_train = xx1[idx, :]\n",
    "u_i_train = uu1[idx, :]\n",
    "\n",
    "# adding boundary and initial points\n",
    "X_u_train = np.vstack([X_u_train, X_i_train])\n",
    "u_train = np.vstack([u_train, u_i_train])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D = Discriminator(in_dim = 3, out_dim = 1, hid_dim = d_hid_dim, num_layers = d_num_layer).to(device)\n",
    "G = Generator(in_dim = 3, out_dim = 1, hid_dim = g_hid_dim, num_layers = g_num_layer).to(device)\n",
    "Q = Q_Net(in_dim = 3, out_dim = 1, hid_dim = q_hid_dim, num_layers = q_num_layer).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "burgers = Burgers_PIG(X_u_train, u_train, X_f_train, X_star, u_star, G, D, Q, device, num_epochs, lambda_phy, noise)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "burgers.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Xmean = burgers.Xmean\n",
    "Xstd = burgers.Xstd\n",
    "X_star_norm = (X_star - Xmean) / Xstd\n",
    "nsamples = 500\n",
    "u_pred_list = []\n",
    "f_pred_list = []\n",
    "for run in range(nsamples):\n",
    "    u_pred, f_pred = burgers.predict(X_star_norm)\n",
    "    u_pred_list.append(u_pred)\n",
    "    f_pred_list.append(f_pred)\n",
    "\n",
    "    \n",
    "u_pred_arr = np.array(u_pred_list)\n",
    "f_pred_arr = np.array(f_pred_list)\n",
    "u_pred = u_pred_arr.mean(axis=0)\n",
    "f_pred = f_pred_arr.mean(axis=0)\n",
    "u_dev = u_pred_arr.var(axis=0)\n",
    "f_dev = f_pred_arr.var(axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "error_u = np.linalg.norm(u_star-u_pred,2)/np.linalg.norm(u_star,2)\n",
    "print('Error u: %e' % (error_u))                     \n",
    "print('Residual: %e' % (f_pred**2).mean())\n",
    "U_pred = griddata(X_star, u_pred.flatten(), (X, T), method='cubic')\n",
    "U_dev = griddata(X_star, u_dev.flatten(), (X, T), method='cubic')\n",
    "Error = np.abs(Exact - U_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" The aesthetic setting has changed. \"\"\"\n",
    "\n",
    "####### Row 0: u(t,x) ##################    \n",
    "X_u_train_ = X_u_train \n",
    "fig = plt.figure(figsize=(9, 5))\n",
    "ax = fig.add_subplot(111)\n",
    "t = data['t'].flatten()[:,None]\n",
    "x = data['x'].flatten()[:,None]\n",
    "\n",
    "h = ax.imshow(Exact.T, interpolation='nearest', cmap='rainbow', \n",
    "              extent=[t.min(), t.max(), x.min(), x.max()], \n",
    "              origin='lower', aspect='auto')\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.10)\n",
    "cbar = fig.colorbar(h, cax=cax)\n",
    "cbar.ax.tick_params(labelsize=15) \n",
    "\n",
    "ax.plot(\n",
    "    X_u_train_[:,1], \n",
    "    X_u_train_[:,0], \n",
    "    'kx', label = 'Data (%d points)' % (u_train.shape[0]), \n",
    "    markersize = 4,  # marker size doubled\n",
    "    clip_on = False,\n",
    "    alpha=1.0\n",
    ")\n",
    "\n",
    "line = np.linspace(x.min(), x.max(), 2)[:,None]\n",
    "ax.plot(t[25]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[50]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[75]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "\n",
    "ax.set_xlabel('$t$', size=20)\n",
    "ax.set_ylabel('$x$', size=20)\n",
    "ax.legend(\n",
    "    loc='upper center', \n",
    "    bbox_to_anchor=(0.9, -0.05), \n",
    "    ncol=5, \n",
    "    frameon=False, \n",
    "    prop={'size': 15}\n",
    ")\n",
    "ax.set_title('$u(t,x)$', fontsize = 20) # font size doubled\n",
    "ax.tick_params(labelsize=15)\n",
    "\n",
    "####### Row 0: u(t,x) ##################    \n",
    "\n",
    "fig = plt.figure(figsize=(9, 5))\n",
    "ax = fig.add_subplot(111)\n",
    "t = data['t'].flatten()[:,None]\n",
    "x = data['x'].flatten()[:,None]\n",
    "\n",
    "h = ax.imshow(U_pred.T, interpolation='nearest', cmap='rainbow', \n",
    "              extent=[t.min(), t.max(), x.min(), x.max()], \n",
    "              origin='lower', aspect='auto')\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.10)\n",
    "cbar = fig.colorbar(h, cax=cax)\n",
    "cbar.ax.tick_params(labelsize=15) \n",
    "\n",
    "ax.plot(\n",
    "    X_u_train_[:,1], \n",
    "    X_u_train_[:,0], \n",
    "    'kx', label = 'Data (%d points)' % (u_train.shape[0]), \n",
    "    markersize = 4,  # marker size doubled\n",
    "    clip_on = False,\n",
    "    alpha=1.0\n",
    ")\n",
    "\n",
    "line = np.linspace(x.min(), x.max(), 2)[:,None]\n",
    "ax.plot(t[25]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[50]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[75]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "\n",
    "ax.set_xlabel('$t$', size=20)\n",
    "ax.set_ylabel('$x$', size=20)\n",
    "ax.legend(\n",
    "    loc='upper center', \n",
    "    bbox_to_anchor=(0.9, -0.05), \n",
    "    ncol=5, \n",
    "    frameon=False, \n",
    "    prop={'size': 15}\n",
    ")\n",
    "ax.set_title('$u(t,x)$', fontsize = 20) # font size doubled\n",
    "ax.tick_params(labelsize=15)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "####### Row 0: u(t,x) ##################    \n",
    "\n",
    "fig = plt.figure(figsize=(9, 5))\n",
    "ax = fig.add_subplot(111)\n",
    "t = data['t'].flatten()[:,None]\n",
    "x = data['x'].flatten()[:,None]\n",
    "\n",
    "h = ax.imshow(U_dev.T, interpolation='nearest', cmap='rainbow', \n",
    "              extent=[t.min(), t.max(), x.min(), x.max()], \n",
    "              origin='lower', aspect='auto')\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.10)\n",
    "cbar = fig.colorbar(h, cax=cax)\n",
    "cbar.ax.tick_params(labelsize=15) \n",
    "\n",
    "ax.plot(\n",
    "    X_u_train_[:,1], \n",
    "    X_u_train_[:,0], \n",
    "    'kx', label = 'Data (%d points)' % (u_train.shape[0]), \n",
    "    markersize = 4,  # marker size doubled\n",
    "    clip_on = False,\n",
    "    alpha=1.0\n",
    ")\n",
    "\n",
    "line = np.linspace(x.min(), x.max(), 2)[:,None]\n",
    "ax.plot(t[25]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[50]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[75]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "\n",
    "ax.set_xlabel('$t$', size=20)\n",
    "ax.set_ylabel('$x$', size=20)\n",
    "ax.legend(\n",
    "    loc='upper center', \n",
    "    bbox_to_anchor=(0.9, -0.05), \n",
    "    ncol=5, \n",
    "    frameon=False, \n",
    "    prop={'size': 15}\n",
    ")\n",
    "ax.set_title('$u(t,x)$', fontsize = 20) # font size doubled\n",
    "ax.tick_params(labelsize=15)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "####### Row 1: u(t,x) slices ##################    \n",
    "gs1 = gridspec.GridSpec(1, 3)\n",
    "gs1.update(top=1-1/3, bottom=0, left=0.1, right=0.9, wspace=0.5)\n",
    "\n",
    "fig = plt.figure(figsize=(20, 10))\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "ax = plt.subplot(gs1[0, 0])\n",
    "ax.plot(x,Exact[25,:], 'b-', linewidth = 2, label = 'Exact')       \n",
    "ax.plot(x,U_pred[25,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "lower = U_pred[25,:] - 2.0*np.sqrt(U_dev[25,:])\n",
    "upper = U_pred[25,:] + 2.0*np.sqrt(U_dev[25,:])\n",
    "plt.fill_between(x.flatten(), lower.flatten(), upper.flatten(), \n",
    "                 facecolor='orange', alpha=0.5, label=\"Two std band\")\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$u(t,x)$')    \n",
    "ax.set_title('$t = 0.25$', fontsize = 10)\n",
    "ax.axis('square')\n",
    "ax.set_xlim([-1.1,1.1])\n",
    "# ax.set_ylim([-1.1,1.1])\n",
    "\n",
    "ax = plt.subplot(gs1[0, 1])\n",
    "ax.plot(x,Exact[50,:], 'b-', linewidth = 2, label = 'Exact')       \n",
    "ax.plot(x,U_pred[50,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "lower = U_pred[50,:] - 2.0*np.sqrt(U_dev[50,:])\n",
    "upper = U_pred[50,:] + 2.0*np.sqrt(U_dev[50,:])\n",
    "plt.fill_between(x.flatten(), lower.flatten(), upper.flatten(), \n",
    "                 facecolor='orange', alpha=0.5, label=\"Two std band\")\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$u(t,x)$')\n",
    "ax.axis('square')\n",
    "ax.set_xlim([-1.1,1.1])\n",
    "# ax.set_ylim([-1.1,1.1])\n",
    "ax.set_title('$t = 0.50$', fontsize = 10)\n",
    "ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.35), ncol=5, frameon=False)\n",
    "\n",
    "ax = plt.subplot(gs1[0, 2])\n",
    "ax.plot(x,Exact[75,:], 'b-', linewidth = 2, label = 'Exact')       \n",
    "ax.plot(x,U_pred[75,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "lower = U_pred[75,:] - 2.0*np.sqrt(U_dev[75,:])\n",
    "upper = U_pred[75,:] + 2.0*np.sqrt(U_dev[75,:])\n",
    "plt.fill_between(x.flatten(), lower.flatten(), upper.flatten(), \n",
    "                 facecolor='orange', alpha=0.5, label=\"Two std band\")\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$u(t,x)$')\n",
    "ax.axis('square')\n",
    "ax.set_xlim([-1.1,1.1])\n",
    "# ax.set_ylim([-1.1,1.1])    \n",
    "ax.set_title('$t = 0.75$', fontsize = 10)\n",
    "\n",
    "\n",
    "fig, ax = newfig(1.0)\n",
    "ax.axis('off')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
